from query_rewrite import DefaultHmmParams, viterbi, simplify_pinyin, get_split_py_for_list, dag, DefaultDagParams
from pypinyin import lazy_pinyin
'''
emission_dict: 存储每个汉字对应的拼音的概率，大部分为单音字，概率是1.0
pinyin2hanzi_dict: 存储每个拼音对应的所有汉字，汉字来连在一起，是一个字符串：一亿益...，用来遍历拼音对应的每个汉字
start_dict: 每个汉字对应的概率，初始汉字使用，因为第一个汉字前面没有汉字，不能使用transition_dict
transition_dict: 存储每个汉字后面跟随的汉字及其概率，数量较多，按照过来从大到小排列
'''
def test_text(model_name) :
    row_query_list = ['感冒零颗力', '地耐的乳膏', '感冒零教攮', '安宫牛黄玩', '本草刚木', '高血压', '消化道细统用药', '维尔唐加瑞', '青梅素',
                      '儿康ningtang浆', '氨咖皇闵胶囊', '阿杞美素胶囊', '疯繆肠胃康分散片', '三久胃太颗粒', '白云山板兰跟颗粒', '白云山富方丹参片', '丫包丁桂儿脐贴',
                      '溏尿病', '新脑血管', '毛细孔情头洁面乳', '儿童相城香波', '黄瓜水润洗发暮雨露', '梅百柔肤水', '杜雷斯', '健安洗', '竹碳黑碳美白面膜',
                      '美帝惠尔', '钻石贵富', '美白补水保湿面馍', '青岛柯美生物工成有限公司', '六炜地皇胶囊', '东啊啊胶', '长春海月药业股份', '葡萄糖算盖口服耶',
                      '感冒灵颗粒', '地奈德乳膏', '感冒灵胶囊']
    row_query_list2 = ['感冒零颗力']
    for row_query in row_query_list2:
        ob = lazy_pinyin(row_query)
        ob_new = get_split_py_for_list([simplify_pinyin(s) for s in ob])
        if model_name == 'hmm':
            hmmparams = DefaultHmmParams()
            result = viterbi(hmm_params=hmmparams, observations=ob_new, path_num=1)
        else:
            dagparams = DefaultDagParams()
            result = dag(dag_params=dagparams, pinyin_list=ob_new, path_num=2)
        for item in result:
            print(row_query + "---->", "".join(item.path))

# 2个候选
if __name__ == '__main__':
    test_text('dag')
    # while True:
    #     row_query = input("请输入原始Query: ")
    #     ob = lazy_pinyin(row_query)
    #     ob_new = get_split_py_for_list([simplify_pinyin(s) for s in ob])
    #     hmmparams = DefaultHmmParams()
    #     result = viterbi(hmm_params=hmmparams, observations=ob_new, path_num=1)
    #     for item in result:
    #         print(row_query + "---->", "".join(item.path))